# SSH server for testing
# based on http://www.devshed.com/c/a/Python/SSH-with-Twisted/

import threading
from twisted.cred import checkers, credentials
from twisted.cred.portal import Portal, IRealm
from twisted.cred.checkers import InMemoryUsernamePasswordDatabaseDontUse
from twisted.conch import error, avatar, recvline, interfaces as conchinterfaces
from twisted.conch.ssh import factory, userauth, connection, keys, session, common 
from twisted.conch.ssh.keys import Key
from twisted.conch.insults import insults 
from twisted.application import service, internet
from twisted.python.randbytes import insecureRandom
from twisted.internet import reactor
from zope.interface import implements


class TestSSHProtocol(recvline.HistoricRecvLine):
    
    def __init__(self, user):
        self.user = user

    def connectionMade(self) : 
        recvline.HistoricRecvLine.connectionMade(self)
        self.terminal.write("Welcome to my test SSH server.")
        self.terminal.nextLine() 
        self.do_help()
        self.showPrompt()

    def showPrompt(self): 
        self.terminal.write("$ ")

    def getCommandFunc(self, cmd):
        return getattr(self, 'do_' + cmd, None)

    def lineReceived(self, line):
        line = line.strip()
        if line: 
            cmdAndArgs = line.split()
            cmd = cmdAndArgs[0]
            args = cmdAndArgs[1:]
            func = self.getCommandFunc(cmd)
            if func: 
               try:
                   func(*args)
               except Exception, e: 
                   self.terminal.write("Error: %s" % e)
                   self.terminal.nextLine()
            else:
               self.terminal.write("No such command.")
               self.terminal.nextLine()
        self.showPrompt()

    def do_help(self, cmd=''):
        "Get help on a command. Usage: help command"
        if cmd: 
            func = self.getCommandFunc(cmd)
            if func:
                self.terminal.write(func.__doc__)
                self.terminal.nextLine()
                return

        publicMethods = filter(
            lambda funcname: funcname.startswith('do_'), dir(self)) 
        commands = [cmd.replace('do_', '', 1) for cmd in publicMethods] 
        self.terminal.write("Commands: " + " ".join(commands))
        self.terminal.nextLine()

    def do_echo(self, *args):
        "Echo a string. Usage: echo my line of text"
        self.terminal.write(" ".join(args)) 
        self.terminal.nextLine()

    def do_whoami(self):
        "Prints your user name. Usage: whoami"
        self.terminal.write(self.user.username)
        self.terminal.nextLine()

    def do_quit(self):
        "Ends your session. Usage: quit" 
        self.terminal.write("Thanks for playing!")
        self.terminal.nextLine() 
        self.terminal.loseConnection()

    def do_clear(self):
        "Clears the screen. Usage: clear" 
        self.terminal.reset()


class TestSSHAvatar(avatar.ConchUser): 
    implements(conchinterfaces.ISession)

    def __init__(self, username): 
        avatar.ConchUser.__init__(self) 
        self.username = username 
        self.channelLookup.update({'session':session.SSHSession})

    def openShell(self, protocol): 
        serverProtocol = insults.ServerProtocol(TestSSHProtocol, self)
        serverProtocol.makeConnection(protocol)
        protocol.makeConnection(session.wrapProtocol(serverProtocol))

    def getPty(self, terminal, windowSize, attrs):
        return None

    def execCommand(self, protocol, cmd): 
        raise NotImplementedError

    def closed(self):
        pass


class TestSSHRealm(object):
    implements(IRealm)

    def requestAvatar(self, avatarId, mind, *interfaces):
        if conchinterfaces.IConchUser in interfaces:
            return interfaces[0], TestSSHAvatar(avatarId), lambda: None
        else:
            raise Exception("No supported interfaces found.")
            

class Checker(InMemoryUsernamePasswordDatabaseDontUse):
    def __init__(self, users):
        InMemoryUsernamePasswordDatabaseDontUse.__init__(self, **users)
        
    def requestAvatarId(self, credentials):
        return InMemoryUsernamePasswordDatabaseDontUse.requestAvatarId(self, credentials)

 
def generateKey():
    print "Generating RSA keypair" 
    from Crypto.PublicKey import RSA 
    KEY_LENGTH = 1024
    return Key(RSA.generate(KEY_LENGTH, insecureRandom))
    

class TestSSHFactory(factory.SSHFactory):
    portal = Portal(TestSSHRealm())
    users = {'admin': 'aaa', 'guest': 'bbb'}
    portal.registerChecker(Checker(users))

    key = generateKey()
    publicKeys = {'ssh-rsa': key}
    privateKeys = {'ssh-rsa': key}


class ReactorThread(threading.Thread):
    def run(self):
        reactor.run(installSignalHandlers=0)
  
def start_ssh_server(port=2222):
    sshFactory = TestSSHFactory() 
    reactor.listenTCP(port, sshFactory)
    thread = ReactorThread()
    print "Starting SSH server"
    thread.start()
    
def stop_ssh_server():
    reactor.stop()
    
if __name__ == "__main__":
    start_ssh_server()
